import argparse


def parse_args():
    parser = argparse.ArgumentParser(description="KGIN")

    # ===== dataset ===== #
    parser.add_argument("--dataset", nargs="?", default="KGIN_data")
    parser.add_argument(
        "--data_path", nargs="?", default="data/", help="Input data path."
    )

    # ===== train ===== #
    parser.add_argument('--model', type=str, default='KGIN_RI')
    parser.add_argument('--version', type=str, default='V1.0', help='the version of model')

    parser.add_argument("--input_num", type=bool, default=True, help="input num(for nodes,entity,items,users) or not")
    parser.add_argument('--n_user', type=int, default=200000, help='number of users')
    parser.add_argument('--n_item', type=int, default=195244, help='number of items. run poi set 29924; run spu set 195244')
    parser.add_argument('--n_entity', type=int, default=198820, help='number of users. run poi set 29924; run spu set 195244; run spu set KGIN 212133')
    parser.add_argument('--n_node', type=int, default=398820, help='number of users. run poi set 229924; run spu set 395244; run spu set KGIN 412133')

    parser.add_argument('--n_relation', type=int, default=2, help='number of relation')

    # 是否对用户点击去重
    # 注意这种bool类型，传过来的都是True，不写加默认才行。
    parser.add_argument("--unique_cf", type=bool, default=False)

    # 控制负采样的种类，all代表从所有物品中采样，recall代表仅仅从召回层中采样
    parser.add_argument('--negative_sampling', type=str, default='all')

    parser.add_argument('--save_model_id', type=int, default=1, help='saved model ckpt name')

    # 控制推荐之前用户点击过的和没有点击过的trade-off
    parser.add_argument('--top_generate_setting', type=str, default='include_train',
                        help='no_train, include_train')

    # 选择是训练模型还是读取之前训练好的模型直接评估HR,设置为evaluate模式后，还会输出给每个用户推荐的物品到文件中
    parser.add_argument('--mode', type=str, default='train',
                        help='train, evaluate')

    # 设置使用的kg_final的文件名
    parser.add_argument('--kg', type=str, default='kg_final_5')

    # 在top_generate_setting设为include_train的时候，这个参数控制添加多少可能复购的菜品
    parser.add_argument('--rebuy_num', type=int, default=50)

    # 设置使用的训练集,测试集
    parser.add_argument('--train_set', type=str, default='train_ori')
    parser.add_argument('--test_set', type=str, default='test', help='test;test_2')

    # 设置使用的召回数据集
    parser.add_argument('--recall_set', type=str, default='recall_items_2w')

    parser.add_argument('--log_comment', type=str, default='lr-1e-2_dim-64', help='comment for log')
    parser.add_argument('--epoch', type=int, default=2, help='number of epochs')
    parser.add_argument('--batch_size', type=int, default=80960, help='batch size')
    parser.add_argument('--test_batch_size', type=int, default=20480, help='batch size')
    parser.add_argument('--dim', type=int, default=64, help='embedding size')
    parser.add_argument('--fm_dim', type=int, default=4, help='fm v size')
    parser.add_argument('--l2', type=float, default=0, help='l2 regularization weight')
    parser.add_argument('--lr', type=float, default=1e-2, help='learning rate')
    parser.add_argument('--sim_regularity', type=float, default=1e-4, help='regularization weight for latent factor')
    parser.add_argument("--inverse_r", type=bool, default=False, help="consider inverse relation or not")
    parser.add_argument("--node_dropout", type=bool, default=True, help="consider node dropout or not")
    parser.add_argument("--node_dropout_rate", type=float, default=0.1, help="ratio of node dropout")
    parser.add_argument("--mess_dropout", type=bool, default=True, help="consider message dropout or not")
    parser.add_argument("--mess_dropout_rate", type=float, default=0.6, help="ratio of node dropout")
    parser.add_argument("--batch_test_flag", type=bool, default=True, help="batch test or not")
    parser.add_argument("--channel", type=int, default=8, help="hidden channels for model, relation embedding size")
    parser.add_argument("--cuda", type=bool, default=True, help="use gpu or not")
    parser.add_argument("--gpu_id", type=int, default=0, help="gpu id")
    parser.add_argument('--Ks', nargs='?', default='[50]', help='Output sizes of every layer') # 论文中的[20, 40, 60, 80, 100],如果用toy数据集要改小一些。
    parser.add_argument('--test_flag', nargs='?', default='part',
                        help='Specify the test type from {part, full}, indicating whether the reference is done in mini-batch')
    parser.add_argument("--n_factors", type=int, default=4, help="number of latent factor for user favour")
    parser.add_argument("--ind", type=str, default='distance', help="Independence modeling: mi, distance, cosine")

    # 设置通过loss或者HR评估模型什么时候训练提前停止。注意：如果设置为loss，那么save参数必须为True，也就是必须训练过程中保存模型。
    parser.add_argument("--stop_train_condition", type=str, default='loss', help="loss, HR")
    # 对于loss停止条件，设置从哪个epoch开始
    parser.add_argument("--loss_start_epoch", type=int, default=0)

    # ===== relation context ===== #
    parser.add_argument('--context_hops', type=int, default=2, help='number of context hops')

    # ===== save model ===== #
    parser.add_argument("--save", type=bool, default=True, help="save model or not")
    parser.add_argument("--out_dir", type=str, default="./weights/", help="output directory for model")

    return parser.parse_args()
